-- Copyright (c) 2024 Huawei Technologies Co., Ltd.
-- openUBMC is licensed under Mulan PSL v2.
-- You can use this software according to the terms and conditions of the Mulan PSL v2.
-- You may obtain a copy of Mulan PSL v2 at: http://license.coscl.org.cn/MulanPSL2
--
-- THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
-- EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
-- MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
-- See the Mulan PSL v2 for more details.
local cjson = require 'cjson'
local m = {}
local MEM_MDB_INTF <const> = "bmc.kepler.Systems.Memory"
local log = require 'mc.logging'
local BIT_ECC_DEFAULT <const> = '{"AggregateTotalCount": 0, "Count": 0, "Info": null}'
local BIT_ISOLATED_DEFAULT <const> ='{"Count": 0, "Info": null}'
local PCIECARD_INTERFACE<const> = 'bmc.kepler.Systems.PCIeDevices.PCIeCard'
local PROPERTY_INTERFACE<const> = 'bmc.kepler.Object.Properties'
local INVALID_PATH<const> = ''

function m.format_identification(processor_id, processor_id_str)
    if processor_id_str then
        return processor_id_str
    end
    if not processor_id then
        return nil
    end
    local str = string.format('%016X', processor_id)
    local identification = ''
    for i = #str, 1, -2 do
        identification = identification .. string.sub(str, i - 1, i)
        if i ~= 2 then
            identification = identification .. '-'
        end
    end
    return identification
end

-- 根据npu对象路径判断其是否为NPUCard
local function get_pcie_card(npu_path, systemid)
    local npu_obj = mdb.get_object(bus, npu_path, PROPERTY_INTERFACE)
    local npu_position = npu_obj.ObjectIdentifier[4]
    local pcie_card_path = "/bmc/kepler/Systems/" .. systemid .. "/PCIeDevices/PCIeCards"
    local pcie_cards = mdb.get_sub_objects(bus, pcie_card_path, PROPERTY_INTERFACE)
    local target_pcie_card
    local ok
    for card_path, card in pairs(pcie_cards) do 
        if card.ObjectIdentifier[4] == npu_position then
            ok, target_pcie_card = pcall(mdb.get_object, bus, card_path, PCIECARD_INTERFACE)
            if ok then 
                return target_pcie_card
            end
        end
    end
end

function m.get_device_physical_id(cpu_paths, systemid)
    local ids = {}
    for _, path in ipairs(cpu_paths) do
        if string.match(path, "CPU") then
            local obj = mdb.get_object(bus, path, 'bmc.kepler.Systems.Processor.CPU')
            table.insert(ids, obj.PhysicalId)
        elseif string.match(path, "GPU") then
            local obj = mdb.get_object(bus, path, 'bmc.kepler.Systems.Processor.GPU')
            table.insert(ids, 'Gpu' .. obj.Slot)
        elseif string.match(path, "NPU") then
            local obj = mdb.get_object(bus, path, 'bmc.kepler.Systems.Processor')
            local ok, res = pcall(get_pcie_card, path, systemid)
            if ok and res then
                table.insert(ids, 'NPU' .. (obj.Id + 1)//2 .. '-' .. ((obj.Id + 1) % 2 + 1))
            else
                table.insert(ids, 'Npu' .. obj.Id)
            end
        end
    end
    return ids
end

function m.get_single_device_physical_id(path)
    if string.match(path, "CPU") then
        local obj = mdb.get_object(bus, path, 'bmc.kepler.Systems.Processor.CPU')
        return obj.PhysicalId
    elseif string.match(path, "GPU") then
        local obj = mdb.get_object(bus, path, 'bmc.kepler.Systems.Processor.GPU')
        return 'Gpu' .. obj.Slot
    elseif string.match(path, "NPU") then
        local obj = mdb.get_object(bus, path, 'bmc.kepler.Systems.Processor')
        return 'Npu' .. obj.Id
    else
        return ''
    end
end

function m.format_nvlink_info(nvlink_info)
    for i = 1, #nvlink_info do
        for j = 1, #nvlink_info[i] do
            if nvlink_info[i][j] == 32768 then
                nvlink_info[i][j] = "NA"
            end
        end
    end
    return nvlink_info
end

function m.get_gpu_nvlink_info(nvlink_info)
    local param_idx = {
        [1] = "NvLinkStatus",
        [2] = "ReplayErrorCount",
        [3] = "RecoveryErrorCount",
        [4] = "FlitCRCErrorCount",
        [5] = "DataCRCErrorCount"
    }
    local info = cjson.json_object_new_array()

    for i = 1, #nvlink_info do
        local element = cjson.json_object_new_object()
        for j = 1, #nvlink_info[i] do
            if nvlink_info[i][j] == 32768 then
                nvlink_info[i][j] = "NA"
            end
            element[param_idx[j]] = nvlink_info[i][j]
        end
        info[i] = element
    end
    return info
end

function m.get_hbm_path(mem_paths, cpuid)
    local res_path = ''
    for _, path in ipairs(mem_paths) do
        local ok, obj = pcall(mdb.get_object, bus, path, MEM_MDB_INTF)
        if ok and obj.CpuId == cpuid and string.find(obj.MemoryDeviceType, 'DRAM') then
            return path
        end
    end
    return res_path
end

function m.get_npu_ecc(ecc_info, info_type)
    if ecc_info then
        return cjson.json_object_ordered_decode(ecc_info)
    end
    log:error('Get %s failed.', info_type)
    if info_type == 'SingleBitEcc' or info_type == 'MultiBitEcc' then
        return cjson.json_object_ordered_decode(BIT_ECC_DEFAULT)
    end
    return cjson.json_object_ordered_decode(BIT_ISOLATED_DEFAULT)
end

function m.get_pcie_npu_device_locator(npu_path, systemid)
    local pcie_card = get_pcie_card(npu_path, systemid)
    if not pcie_card then
        return cjson.null
    end
    return 'PCIe Card ' .. pcie_card.SlotID .. ' (NPU)'
end

function m.get_pcie_npu_position(npu_path, systemid)
    local pcie_card = get_pcie_card(npu_path, systemid)
    if not pcie_card then
        return cjson.null
    end
    local container_slot = string.match(pcie_card.Position, '%d+')
    return 'PCIe Riser' .. container_slot
end

return m